package edu.vt.cbil.dataset;

/**
 * @author Justin
 * @version 1.0
 * @created 08-11月-2014 22:18:57
 */

import edu.vt.cbil.util.BioToolbox;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;

public class AnnotationDataSet {

    private Map<String, AnnotationProbeSet> ProbeSets;
    private String AnnotationFilePath;

    public AnnotationDataSet(String filename) {
        ProbeSets = new HashMap<String, AnnotationProbeSet>(180 * 10000);
        AnnotationFilePath = filename;
    }


    public Map<String, AnnotationProbeSet> getProbeSets() {
        return ProbeSets;
    }

    /**
     * @param ProbeSetID: the input loc
     * @return : the ATCG code of Allele A at the input loc
     */
    public char getAlleleACode(String ProbeSetID) {
        return ProbeSets.get(ProbeSetID).Allele_A;
    }

    public char getAlleleBCode(String ProbeSetID) {
        return ProbeSets.get(ProbeSetID).Allele_B;
    }

    /**
     * @param ProbeSetID
     * @return the annotation data of the input loc
     */
    public AnnotationProbeSet get(String ProbeSetID) {
        return ProbeSets.get(ProbeSetID);
    }

    public void readingAnnotation(boolean onlySNP) {

        try {
            // connect sqlite database
            Class.forName("org.sqlite.JDBC");
            String ConnectionString = String.format("jdbc:sqlite:%s", AnnotationFilePath);
            Connection c = DriverManager.getConnection(ConnectionString);

            Statement stmt = c.createStatement();
            ResultSet rs;
            if (onlySNP) {
                rs = stmt.executeQuery("SELECT ProbeSet_ID, Chr_id,Start,Stop,Allele_A,Allele_B, Cytoband FROM "
                        + " Annotations WHERE ProbeSet_ID LIKE 'S%' ORDER BY Chr_id, Start;");
            } else {
                rs = stmt.executeQuery("SELECT ProbeSet_ID, Chr_id,Start,Stop,Allele_A,Allele_B, Cytoband FROM "
                        + " Annotations ORDER BY Chr_id, Start;");
            }

            ProbeSets.clear();

            while (rs.next()) {

                AnnotationProbeSet aps = new AnnotationProbeSet();

                aps.ProbeSetID = rs.getString("ProbeSet_ID").trim();
                aps.Chr_id = rs.getShort("Chr_id");
                aps.Start = rs.getInt("Start");
                aps.Stop = rs.getInt("Stop");
                aps.Cytoband = rs.getString("Cytoband").trim();

                aps.ProbeType = (aps.ProbeSetID.startsWith("CN")) ? ProbeSetType.CN
                        : ProbeSetType.SNP;

                if (onlySNP && (aps.ProbeType == ProbeSetType.CN)) {
                    continue;
                }

                if (aps.ProbeType == ProbeSetType.CN) {
                    //aps.Pos = ((long)aps.Chr_id << 480) | (long)((aps.Start + aps.Stop)/2); //使用中间点
                    aps.Pos = BioToolbox.getGlobalLocation(aps.Chr_id, (aps.Start + aps.Stop) / 2);
                } else {
                    //aps.Pos = ((long)aps.Chr_id << 480) | (long)aps.Start;
                    aps.Pos = BioToolbox.getGlobalLocation(aps.Chr_id, aps.Start);
                }

                if (aps.ProbeType == ProbeSetType.SNP) {
                    aps.Allele_A = rs.getString("Allele_A").trim().toUpperCase().charAt(0);
                    aps.Allele_B = rs.getString("Allele_B").trim().toUpperCase().charAt(0);
                }

                ProbeSets.put(aps.ProbeSetID, aps);
            }
            //System.out.println(ProbeSets.size());
            rs.close();
            stmt.close();
            c.close();

        } catch (Exception e) {
            System.err.println(e.getClass().getName() + ": " + e.getMessage());
            System.exit(0);
        }

    }

}